package dd.lo.chapter12;

import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.BorderPane;
import javafx.scene.layout.FlowPane;
import javafx.scene.text.Text;
import javafx.stage.FileChooser;
import javafx.stage.Stage;
import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.util.*;

/**
 * Watershed algorithm example
 * 参考资料：
 * https://blog.csdn.net/Vichael_Chan/article/details/100931620
 * https://blog.csdn.net/liyuanbhu/article/details/119382070
 */
public class Example4 extends Application {

    public static void main(String[] args) {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);
        launch(args);
    }

    private static final int DEFAULT_WIDTH = 1080;
    private static final int DEFAULT_HEIGHT = 768;

    private enum Mode {
        ORIGIN,
        GRAY,
        GAUSSIAN_BLUR,
        CANNY,
        CONTOUR,
        MARK,
        WATERSHED,
        COLORED,
        BLENDED,
    }

    private int modeIdx = 0;

    private Mat srcImg;

    private Mat grayImg;

    private Mat gaussianBlurredImg;

    private Mat cannyImg;

    private Mat contourImg;

    private Mat markImg;

    private Mat watershedImg;

    private Mat coloredImg;

    private Mat blendedImg;

    private ImageView imageView;

    private Text modeText;

    private boolean calculating = false;

    private Text stateText;

    private double[] getRandomColor(int value) {
        value %= 255;
        double r = (Math.random() * value);
        double g = (Math.random() * value);
        double b = (Math.random() * value);
        return new double[]{r, g, b};
    }

    private void imageProcessing() {
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                if (null == srcImg) return null;
                calculating = true;
                Platform.runLater(() -> stateText.setText(calculating ? "Calculating" : "Ready"));
                srcImg.copyTo(watershedImg = new Mat());
                Imgproc.cvtColor(srcImg, grayImg = new Mat(), Imgproc.COLOR_BGR2GRAY);
                Imgproc.GaussianBlur(grayImg, gaussianBlurredImg = new Mat(), new Size(5, 5), 2);
                Imgproc.Canny(gaussianBlurredImg, cannyImg = new Mat(), 80, 150);
                List<MatOfPoint> contours = new ArrayList<>();
                Mat hierarchy = new Mat();
                Imgproc.findContours(cannyImg, contours, hierarchy, Imgproc.RETR_TREE, Imgproc.CHAIN_APPROX_SIMPLE);
                contourImg = Mat.zeros(srcImg.size(), CvType.CV_8UC1);
                markImg = new Mat(srcImg.size(), CvType.CV_32S); //OpenCv分水岭第二个参数
                for (int i = 0; i < contours.size(); ++i) {
                    Imgproc.drawContours(markImg, contours, i, Scalar.all(i + 1), 1, Imgproc.LINE_AA, hierarchy);
                    Imgproc.drawContours(contourImg, contours, i, Scalar.all(255), 1, Imgproc.LINE_AA, hierarchy);
                }
                Imgproc.watershed(watershedImg, markImg);
                Map<Integer, double[]> valueMapColor = new HashMap<>();
                coloredImg = Mat.zeros(srcImg.size(), CvType.CV_8UC3);
                for (int i = 0; i < markImg.rows(); ++i) {
                    for (int j = 0; j < markImg.cols(); ++j) {
                        int index = (int) markImg.get(i, j)[0];
                        if (index == -1) {
                            coloredImg.put(i, j, 255, 255, 255);
                        } else {
                            double[] color = valueMapColor.get(index);
                            if (null == color) {
                                color = getRandomColor(index);
                                valueMapColor.put(index, color);
                            }
                            coloredImg.put(i, j, color);
                        }
                    }
                }
                Core.addWeighted(srcImg, 0.4, coloredImg, 0.6, 0, blendedImg = new Mat());
                Core.convertScaleAbs(markImg, watershedImg);
                calculating = false;
                Platform.runLater(() -> stateText.setText(calculating ? "Calculating" : "Ready"));
                return null;
            }
        };
        new Thread(task).start();
    }

    private void imageShow(Mode mode) {
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                Mat dst = new Mat();
                switch (mode) {
                    case ORIGIN: {
                        srcImg.copyTo(dst);
                        break;
                    }
                    case GRAY: {
                        grayImg.copyTo(dst);
                        break;
                    }
                    case GAUSSIAN_BLUR: {
                        gaussianBlurredImg.copyTo(dst);
                        break;
                    }
                    case CANNY: {
                        cannyImg.copyTo(dst);
                        break;
                    }
                    case CONTOUR: {
                        contourImg.copyTo(dst);
                        break;
                    }
                    case MARK: {
                        markImg.copyTo(dst);
                        break;
                    }
                    case WATERSHED: {
                        watershedImg.copyTo(dst);
                        break;
                    }
                    case COLORED: {
                        coloredImg.copyTo(dst);
                        break;
                    }
                    case BLENDED: {
                        blendedImg.copyTo(dst);
                        break;
                    }
                    default: {
                        System.out.printf("无效mode：%s\n", mode);
                    }
                }
                if (dst.empty()) return null;
                //调整图片大小以用于展示
                ImageUtils.fitImgSize(dst, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                MatOfByte buffer = new MatOfByte();
                Imgcodecs.imencode(".jpeg", dst, buffer);
                Platform.runLater(() -> imageView.setImage(new Image(new ByteArrayInputStream(buffer.toArray()))));
                return null;
            }
        };
        new Thread(task).start();
    }

    @Override
    public void start(Stage stage) {
        //File chooser
        Button imageChooser = new Button("选择图片");
        imageChooser.setOnMouseClicked(e -> {
            if (calculating) return;
            FileChooser fileChooser = new FileChooser();
            fileChooser.setTitle("请选择图片");
            fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
            );
            fileChooser.getExtensionFilters().addAll(
//                    new FileChooser.ExtensionFilter("All Images", "*.*"),
                    new FileChooser.ExtensionFilter("JPEG", "*.jpeg"),
                    new FileChooser.ExtensionFilter("WEBP", "*.webp"),
                    new FileChooser.ExtensionFilter("PNG", "*.png"),
                    new FileChooser.ExtensionFilter("JPG", "*.jpg"),
                    new FileChooser.ExtensionFilter("GIF", "*.gif"),
                    new FileChooser.ExtensionFilter("BMP", "*.bmp")
            );
            File imgFile = fileChooser.showOpenDialog(stage);
            if (null == imgFile) {
                System.out.println("用户取消选中文件");
                return;
            }
            //读取一张图片
            srcImg = Imgcodecs.imread(imgFile.getAbsolutePath(), Imgcodecs.IMREAD_UNCHANGED);
            if (srcImg.empty()) {
                throw new RuntimeException("图片读取失败");
            }
            modeIdx = 0;
            Platform.runLater(() -> modeText.setText(String.format("Mode=%s", Mode.values()[modeIdx])));
            imageShow(Mode.ORIGIN);
            imageProcessing();
        });
        //Curr mode text
        modeText = new Text(String.format("Mode=%s", Mode.values()[modeIdx]));
        //Prev mode button
        Button prevModeButton = new Button("前一个");
        prevModeButton.setOnMouseClicked(e -> {
            if (calculating) return;
            --modeIdx;
            if (modeIdx < 0) modeIdx = Mode.values().length - 1;
            Platform.runLater(() -> modeText.setText(String.format("Mode=%s", Mode.values()[modeIdx])));
            imageShow(Mode.values()[modeIdx]);
        });
        //Next mode button
        Button nextModeButton = new Button("下一个");
        nextModeButton.setOnMouseClicked(e -> {
            if (calculating) return;
            ++modeIdx;
            if (modeIdx >= Mode.values().length) modeIdx = 0;
            Platform.runLater(() -> modeText.setText(String.format("Mode=%s", Mode.values()[modeIdx])));
            imageShow(Mode.values()[modeIdx]);
        });
        //State text
        stateText = new Text(calculating ? "Calculating..." : "Ready");
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(imageChooser, prevModeButton, modeText, nextModeButton, stateText);
        bottomControl.setPrefHeight(50);
        BorderPane root = new BorderPane();
        imageView = new ImageView();
        root.setCenter(imageView);
        root.setBottom(bottomControl);
        Scene scene = new Scene(root, DEFAULT_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("Example4 - Auto segmenting image via Watershed Algorithm");
        stage.setResizable(false);
        stage.show();
    }
}
